//
// Created for realtime detection pipeline
//

#include "ros/ros.h"
#include <sensor_msgs/Image.h>
#include <std_msgs/String.h>
#include <cv_bridge/cv_bridge.h>
#include <sensor_msgs/image_encodings.h>
#include <message_filters/subscriber.h>
#include <message_filters/time_synchronizer.h>
#include <message_filters/sync_policies/approximate_time.h>
#include <opencv2/opencv.hpp>
#include <memory>
#include <string>

// 自定义ImageData消息类型 (简化版本，也可以定义为ROS消息类型)
struct ImageDataMsg {
    uint64_t frame_id;
    ros::Time timestamp;
    sensor_msgs::Image color_image;
    sensor_msgs::Image depth_image;
    std::string camera_info;  // 可选：相机内参信息
};

class ImageDataPublisher {
private:
    ros::NodeHandle nh_;
    
    // 同步订阅器
    message_filters::Subscriber<sensor_msgs::Image> color_sub_;
    message_filters::Subscriber<sensor_msgs::Image> depth_sub_;
    
    // 时间同步策略
    typedef message_filters::sync_policies::ApproximateTime<sensor_msgs::Image, sensor_msgs::Image> SyncPolicy;
    typedef message_filters::Synchronizer<SyncPolicy> Synchronizer;
    std::shared_ptr<Synchronizer> sync_;
    
    // 发布器
    ros::Publisher color_pub_;
    ros::Publisher depth_pub_;
    ros::Publisher sync_info_pub_;
    
    // 参数
    std::string color_input_topic_;
    std::string depth_input_topic_;
    std::string color_output_topic_;
    std::string depth_output_topic_;
    std::string sync_info_topic_;
    
    // 统计信息
    uint64_t frame_counter_;
    ros::Time last_sync_time_;
    double sync_rate_;
    
    // 显示窗口开关
    bool show_images_;
    bool enable_sync_display_;

public:
    ImageDataPublisher() : 
        nh_("~"),
        color_sub_(nh_, "", 1),
        depth_sub_(nh_, "", 1),
        frame_counter_(0),
        sync_rate_(0.0) {
        
        // 读取参数
        loadParameters();
        
        // 初始化订阅器
        color_sub_.subscribe(nh_, color_input_topic_, 1);
        depth_sub_.subscribe(nh_, depth_input_topic_, 1);
        
        // 创建同步器
        sync_ = std::make_shared<Synchronizer>(SyncPolicy(10), color_sub_, depth_sub_);
        sync_->registerCallback(boost::bind(&ImageDataPublisher::syncCallback, this, _1, _2));
        
        // 初始化发布器
        color_pub_ = nh_.advertise<sensor_msgs::Image>(color_output_topic_, 1);
        depth_pub_ = nh_.advertise<sensor_msgs::Image>(depth_output_topic_, 1);
        sync_info_pub_ = nh_.advertise<std_msgs::String>(sync_info_topic_, 1);
        
        ROS_INFO("=== Image Data Publisher Initialized ===");
        ROS_INFO("Color input: %s", color_input_topic_.c_str());
        ROS_INFO("Depth input: %s", depth_input_topic_.c_str());
        ROS_INFO("Color output: %s", color_output_topic_.c_str());
        ROS_INFO("Depth output: %s", depth_output_topic_.c_str());
        ROS_INFO("Sync info output: %s", sync_info_topic_.c_str());
        ROS_INFO("Show images: %s", show_images_ ? "true" : "false");
        ROS_INFO("========================================");
        
        last_sync_time_ = ros::Time::now();
    }
    
    ~ImageDataPublisher() {
        if (show_images_) {
            cv::destroyAllWindows();
        }
    }

private:
    void loadParameters() {
        // 输入话题
        nh_.param<std::string>("color_input_topic", color_input_topic_, "/camera/color/image_raw");
        nh_.param<std::string>("depth_input_topic", depth_input_topic_, "/camera/depth/image_raw");
        
        // 输出话题
        nh_.param<std::string>("color_output_topic", color_output_topic_, "/image_data/color");
        nh_.param<std::string>("depth_output_topic", depth_output_topic_, "/image_data/depth");
        nh_.param<std::string>("sync_info_topic", sync_info_topic_, "/image_data/sync_info");
        
        // 显示选项
        nh_.param<bool>("show_images", show_images_, true);
        nh_.param<bool>("enable_sync_display", enable_sync_display_, true);
    }
    
    void syncCallback(const sensor_msgs::ImageConstPtr& color_msg,
                     const sensor_msgs::ImageConstPtr& depth_msg) {
        
        // 更新帧计数和同步率统计
        ++frame_counter_;
        updateSyncRate();
        
        try {
            // 转换图像用于显示和处理
            cv_bridge::CvImagePtr color_ptr = cv_bridge::toCvCopy(color_msg, sensor_msgs::image_encodings::BGR8);
            cv_bridge::CvImagePtr depth_ptr;
            
            // 处理不同的深度图像编码
            if (depth_msg->encoding == sensor_msgs::image_encodings::TYPE_16UC1) {
                depth_ptr = cv_bridge::toCvCopy(depth_msg, sensor_msgs::image_encodings::TYPE_16UC1);
            } else if (depth_msg->encoding == sensor_msgs::image_encodings::TYPE_32FC1) {
                depth_ptr = cv_bridge::toCvCopy(depth_msg, sensor_msgs::image_encodings::TYPE_32FC1);
            } else {
                ROS_WARN("Unsupported depth encoding: %s", depth_msg->encoding.c_str());
                return;
            }
            
            // 创建ImageData结构
            ImageDataMsg image_data;
            image_data.frame_id = frame_counter_;
            image_data.timestamp = color_msg->header.stamp;  // 使用彩色图像的时间戳
            image_data.color_image = *color_msg;
            image_data.depth_image = *depth_msg;
            
            // 添加同步信息到header
            sensor_msgs::Image sync_color_msg = *color_msg;
            sensor_msgs::Image sync_depth_msg = *depth_msg;
            
            // 在header中添加frame_id信息
            sync_color_msg.header.frame_id = std::to_string(frame_counter_);
            sync_depth_msg.header.frame_id = std::to_string(frame_counter_);
            
            // 发布同步后的图像数据
            color_pub_.publish(sync_color_msg);
            depth_pub_.publish(sync_depth_msg);
            
            // 发布同步信息
            publishSyncInfo(image_data, color_ptr->image, depth_ptr->image);
            
            // 显示图像
            if (show_images_) {
                displayImages(color_ptr->image, depth_ptr->image, image_data);
            }
            
            // 调试信息
            ROS_DEBUG("Synchronized frame %lu: Color(%dx%d) + Depth(%dx%d), sync_rate: %.1f Hz",
                     frame_counter_,
                     color_ptr->image.cols, color_ptr->image.rows,
                     depth_ptr->image.cols, depth_ptr->image.rows,
                     sync_rate_);
                     
        } catch (cv_bridge::Exception& e) {
            ROS_ERROR("cv_bridge exception: %s", e.what());
        }
    }
    
    void updateSyncRate() {
        ros::Time current_time = ros::Time::now();
        double dt = (current_time - last_sync_time_).toSec();
        
        if (dt > 0) {
            // 简单的滑动平均计算同步率
            double current_rate = 1.0 / dt;
            sync_rate_ = (sync_rate_ * 0.9) + (current_rate * 0.1);
        }
        
        last_sync_time_ = current_time;
    }
    
    void publishSyncInfo(const ImageDataMsg& image_data, 
                        const cv::Mat& color_image, 
                        const cv::Mat& depth_image) {
        std_msgs::String sync_info;
        
        // 创建JSON格式的同步信息
        std::ostringstream oss;
        oss << "{"
            << "\"frame_id\": " << image_data.frame_id << ", "
            << "\"timestamp\": " << image_data.timestamp.toSec() << ", "
            << "\"color_size\": [" << color_image.cols << ", " << color_image.rows << "], "
            << "\"depth_size\": [" << depth_image.cols << ", " << depth_image.rows << "], "
            << "\"sync_rate\": " << sync_rate_ << ", "
            << "\"color_encoding\": \"" << image_data.color_image.encoding << "\", "
            << "\"depth_encoding\": \"" << image_data.depth_image.encoding << "\""
            << "}";
        
        sync_info.data = oss.str();
        sync_info_pub_.publish(sync_info);
    }
    
    void displayImages(const cv::Mat& color_image, const cv::Mat& depth_image, const ImageDataMsg& image_data) {
        // 显示彩色图像
        cv::Mat display_color = color_image.clone();
        
        // 在彩色图像上添加信息
        std::string color_info = cv::format("Frame: %lu, Rate: %.1fHz", 
                                           image_data.frame_id, sync_rate_);
        cv::putText(display_color, color_info, cv::Point(10, 30),
                   cv::FONT_HERSHEY_SIMPLEX, 0.8, cv::Scalar(0, 255, 0), 2);
        
        std::string size_info = cv::format("Color: %dx%d", color_image.cols, color_image.rows);
        cv::putText(display_color, size_info, cv::Point(10, 60),
                   cv::FONT_HERSHEY_SIMPLEX, 0.6, cv::Scalar(0, 255, 255), 1);
        
        cv::imshow("Synchronized Color Image", display_color);
        
        // 显示深度图像
        cv::Mat depth_display;
        if (depth_image.type() == CV_16UC1) {
            // 16位深度图转换为可视化
            depth_image.convertTo(depth_display, CV_8UC1, 255.0/5000.0); // 假设最大深度5米
        } else if (depth_image.type() == CV_32FC1) {
            // 32位浮点深度图转换
            depth_image.convertTo(depth_display, CV_8UC1, 255.0/5.0); // 假设最大深度5米
        } else {
            depth_display = depth_image.clone();
        }
        
        // 应用颜色映射使深度图更直观
        cv::Mat depth_colored;
        cv::applyColorMap(depth_display, depth_colored, cv::COLORMAP_JET);
        
        // 在深度图上添加信息
        std::string depth_info = cv::format("Depth: %dx%d, Type: %s", 
                                           depth_image.cols, depth_image.rows,
                                           image_data.depth_image.encoding.c_str());
        cv::putText(depth_colored, depth_info, cv::Point(10, 30),
                   cv::FONT_HERSHEY_SIMPLEX, 0.6, cv::Scalar(255, 255, 255), 1);
        
        cv::imshow("Synchronized Depth Image", depth_colored);
        
        // 如果启用，显示同步可视化
        if (enable_sync_display_) {
            displaySyncVisualization(color_image, depth_colored, image_data);
        }
        
        cv::waitKey(1);
    }
    
    void displaySyncVisualization(const cv::Mat& color, const cv::Mat& depth, const ImageDataMsg& image_data) {
        // 创建并排显示
        cv::Mat combined;
        cv::Size target_size(640, 480);  // 标准显示尺寸
        
        cv::Mat color_resized, depth_resized;
        cv::resize(color, color_resized, target_size);
        cv::resize(depth, depth_resized, target_size);
        
        cv::hconcat(color_resized, depth_resized, combined);
        
        // 添加整体信息
        std::string sync_info = cv::format("Synchronized Frame %lu - Rate: %.1f Hz", 
                                          image_data.frame_id, sync_rate_);
        cv::putText(combined, sync_info, cv::Point(10, combined.rows - 20),
                   cv::FONT_HERSHEY_SIMPLEX, 0.8, cv::Scalar(0, 255, 0), 2);
        
        cv::imshow("Synchronized RGB-D Data", combined);
    }

public:
    void spin() {
        ros::Rate rate(100);  // 高频率检查，实际同步率由相机决定
        
        ROS_INFO("Starting image data synchronization...");
        ROS_INFO("Waiting for synchronized RGB-D data...");
        
        while (ros::ok()) {
            ros::spinOnce();
            
            // 定期输出统计信息
            if (frame_counter_ > 0 && frame_counter_ % 100 == 0) {
                ROS_INFO("Processed %lu synchronized frames, current rate: %.1f Hz", 
                         frame_counter_, sync_rate_);
            }
            
            rate.sleep();
        }
    }
};

int main(int argc, char **argv) {
    // ROS节点初始化
    ros::init(argc, argv, "get_image_data_cpp");
    
    try {
        ImageDataPublisher publisher;
        publisher.spin();
    } catch (const std::exception& e) {
        ROS_ERROR("Exception in image data publisher: %s", e.what());
        return -1;
    }
    
    return 0;
}